package com.inspur.edp.ai.flow.core.engine.eventbus;

import com.inspur.edp.ai.flow.core.engine.ThreadPoolService;
import com.inspur.edp.ai.flow.engine.EventBus;
import com.inspur.edp.ai.flow.engine.MessageType;
import com.inspur.edp.ai.flow.exception.FlowException;
import com.inspur.edp.ai.flow.model.Output;
import lombok.extern.slf4j.Slf4j;

import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.Consumer;

@Slf4j
public class TopicEventBus implements EventBus {
    private final ConcurrentMap<String, List<Consumer<Output>>> subscribers = new ConcurrentHashMap<>();

    public TopicEventBus() {
    }

    /**
     * 订阅
     *
     * @param topic    主题
     * @param listener 监听者
     */
    @Override
    public void subscribe(String topic, Consumer<Output> listener) {
        subscribers.computeIfAbsent(topic, k -> new CopyOnWriteArrayList<>())
                .add(listener);
    }

    /**
     * 取消监听
     *
     * @param topic    主题
     * @param listener 监听者
     */
    @Override
    public void unsubscribe(String topic, Consumer<Output> listener) {
        if (topic == null) {
            throw new FlowException("订阅主题不能为空");
        }

        List<Consumer<Output>> listeners = subscribers.get(topic);
        if (listeners != null) {
            listeners.remove(listener);
            if (listeners.isEmpty()) {
                subscribers.remove(topic, listeners);
            }
        }
    }

    /**
     * 取消订阅主题
     *
     * @param topic
     * @return
     */
    @Override
    public int unsubscribe(String topic) {
        if (topic == null) throw new FlowException("订阅主题不能为空");
        List<Consumer<Output>> listeners = subscribers.get(topic);
        if (listeners == null) return 0;
        // 对当前 topic 的监听器列表加锁
        synchronized (listeners) {
            int count = listeners.size();
            publishAsync(topic, new Output().setFinishReason(MessageType.STOP.getValue()).setText(""));
            subscribers.remove(topic);
            return count;
        }
    }

    /**
     * @param topic
     * @param output
     */

    @Override
    public void publishAsync(String topic, Output output) {
        List<Consumer<Output>> listeners = subscribers.get(topic);
        if (listeners == null || listeners.isEmpty()) {
            return;
        }
        for (Consumer<Output> listener : listeners) {
            ThreadPoolService.getExecutor().execute(() -> {
                try {
                    listener.accept(output);
                } catch (Throwable e) {
                    log.error("处理事件失败: topic=" + topic + ", event=" + output + ", listener=" + listener, e);
                }
            });
        }
    }

}

